load("@com_github_grpc_grpc//bazel:python_rules.bzl", "py_grpc_library")
load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
load("@rules_proto//proto:defs.bzl", "proto_library")
load("//bazel:api.bzl", "modular_py_library", "requirement")

package(default_visibility = ["//max:consumers"])

modular_py_library(
    name = "kvcache_agent",
    srcs = [
        "__init__.py",
        "dispatcher_v2.py",
        "kvcache_agent.py",
    ],
    imports = ["../../.."],
    deps = [
        requirement("grpcio"),  # NOTE: Order matters for PYTHONPATH nonsense
        "//max/python/max/pipelines/core",
        "//max/python/max/serve:process_control",
        "//max/python/max/serve/queue",
        requirement("pyzmq"),
        requirement("msgspec"),
        ":kvcache_agent_proto_srcs",
    ],
)

modular_py_library(
    name = "kvcache_agent_proto_srcs",
    srcs = [
        ":kvcache_agent_service_v1_grpc_py",
        ":kvcache_agent_service_v1_py",
    ],
    imports = ["../../.."],
    tags = ["no-mypy"],
    deps = [
        requirement("protobuf"),
    ],
)

proto_library(
    name = "kvcache_agent_service_v1",
    srcs = [
        "kvcache_agent_service_v1.proto",
    ],
    visibility = ["//visibility:private"],
)

py_proto_library(
    name = "kvcache_agent_service_v1_py",
    visibility = ["//visibility:private"],
    deps = [":kvcache_agent_service_v1"],
)

py_grpc_library(
    name = "kvcache_agent_service_v1_grpc_py",
    srcs = [":kvcache_agent_service_v1"],
    # By default, the generated stubs use the fully-qualified names
    # We need to strip the prefix from the generated stubs, so that
    # it works with the modular_py_library and the imports above.
    strip_prefixes = ["max.python."],
    visibility = ["//visibility:private"],
    deps = [":kvcache_agent_service_v1_py"],
)
